#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;

template<typename TYPE_X, typename TYPE_Y>
class KernelAsStrided {
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;

    GlobalTensor<TYPE_X> xGm;
    GlobalTensor<TYPE_Y> yGm;
    GM_ADDR x_addr;
    
    uint32_t ALIGN_NUM;

    // uint32_t* shape;
    __gm__ int32_t* viewShape;
    __gm__ int32_t* strideNum;
    uint32_t dimNum;
    uint32_t viewDim;
    int32_t storageOffset;


    uint32_t blockLength;   //总数据aligned
    uint32_t blockSize; //一次tiling最大值

    uint32_t stepSize; //
    uint32_t totalLength; //最后一个维度的长度
    uint32_t unitCount; //前面所有维度积
    // int32_t* unitCountOffset;
    
    
    uint32_t tileNum; //每组reduce数据计算的循环次数
    uint32_t tileLength; //每次tiling的数据长度
    uint32_t lasttileLength; //最后一次tiling的长度/ uint32_t remain_size;

private:
    __aicore__ inline int32_t GetOffset(int index)   //global index
    {
        int32_t product = 1;
        int32_t localid;
        int32_t offset = 0;
        for(int k = viewDim - 2; k > -1; k--)
        {
            localid = index / product % viewShape[k];
            offset += localid * strideNum[k];
            if(k > 0)
                product *= viewShape[k];
        }
        return offset;
    }


public:
    __aicore__ inline KernelAsStrided() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR size, GM_ADDR stride, GM_ADDR storage_offset, GM_ADDR y, 
                                uint32_t blockLength, uint32_t blockSize, uint32_t ALIGN_NUM, uint32_t* shape, uint32_t dimNum, uint32_t viewDim) {
    
    this->blockLength = blockLength;
    this->blockSize = blockSize;
    this->ALIGN_NUM = ALIGN_NUM;

    this->totalLength = 0;

    this->dimNum = dimNum;
    this->viewDim = viewDim;    //视图的维度

    // this->shape = shape;    //输入的维度信息

    this->viewShape = (__gm__ int32_t*)size;  //(3, 3, 3)
    this->strideNum = (__gm__ int32_t*)stride;    //(2, 2, 1)
    this->storageOffset = (*(__gm__ int32_t*)storage_offset);

    this->blockLength = this->blockLength + (this->blockLength % ALIGN_NUM ? ALIGN_NUM - this->blockLength % ALIGN_NUM : 0);    //32Byte对齐

    this->stepSize = strideNum[viewDim-1];
    this->totalLength = viewShape[viewDim-1];
    this->unitCount = 1;

    uint32_t totalLengthAligned = (totalLength + ALIGN_NUM - 1) / ALIGN_NUM * ALIGN_NUM;

    this->tileNum = totalLengthAligned / this->blockSize;
    // 
    if(totalLengthAligned % blockSize == 0 || tileNum == 0) //reduce dim对齐之后，可以均分到各个tile，或不足一个tile
    {
        if(tileNum == 0)    //一次tiling的数据个数 大于 totalLengthAligned，也就是tile_num = 0时，tileLength=totalLengthAligned
        {
            this->tileNum = 1;
            this->tileLength = totalLengthAligned;
            this->lasttileLength = this->totalLength;
        }
        else //一次tiling的数据个数 小于 totalLengthAligned时，tileLength取最大值tiling_size * ALIGN_NUM
        {
            this->tileLength = this->blockSize;
            this->lasttileLength = this->tileLength;
        }
        
    }
    else //totalLengthAligned对齐后不能均分到每次tiling
    {
        this->tileNum = this->tileNum + 1;
        this->tileLength = this->blockSize; //tileLength 取最大值(ub)
        // this->lasttileLength = this->totalLength % this->tileLength;
        this->lasttileLength = totalLengthAligned % this->tileLength;
    }

    ASSERT(tileNum != 0 && "tile num can not be zero!");

    for(int i=0; i<viewDim-1; i++)
    {
        this->unitCount *= viewShape[i];
    }

    
    
    // printf("blockLength: %u \n", blockLength);
    // printf("tileNum: %u \n", tileNum);
    // printf("tileLength: %u \n", tileLength);
    // printf("lasttileLength: %u \n", lasttileLength);
    // printf("storageOffset: %d \n", storageOffset);
    // printf("strideNum: %u \n", strideNum[0]);
    // printf("viewShape: %u \n", viewShape[0]);
    
    // printf("stepSize: %u \n", stepSize);
    // printf("unitCount: %u \n", unitCount);
    // printf("totalLength: %u \n", totalLength);

    this->x_addr = x;
    xGm.SetGlobalBuffer((__gm__ TYPE_X *)x);
    yGm.SetGlobalBuffer((__gm__ TYPE_Y *)y);
    pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(TYPE_X));
    pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileLength * sizeof(TYPE_X));

    }

    __aicore__ inline void Process()
    {
        if(stepSize != 1)
        {
            ProcessV1();
            return;
        }
        //stepSize == 1
        uint32_t loopCount = unitCount;
        for(int i=0; i<loopCount; i++)
        {
            for(int j=0; j<tileNum - 1; j++)
            {
                CopyIn(i, j, tileLength);
                Compute(i, j, tileLength);
                CopyOut(i, j, tileLength);
            }
                CopyIn(i, tileNum - 1, lasttileLength);
                Compute(i, tileNum - 1, lasttileLength);
                CopyOut(i, tileNum - 1, lasttileLength);
        }
    }

    __aicore__ inline void CopyIn(int i, int32_t progress, uint32_t length)
    {
        int32_t startIndex = GetOffset(i) + progress * tileLength + storageOffset;
        LocalTensor<TYPE_X> xLocal = inQueueX.AllocTensor<TYPE_X>();
        DataCopy(xLocal, xGm[startIndex], length);
        inQueueX.EnQue<TYPE_X>(xLocal);
    }

    __aicore__ inline void Compute(int i, int32_t progress, uint32_t length)
    {
        LocalTensor<TYPE_Y> yLocal = outQueueY.AllocTensor<TYPE_Y>();
        LocalTensor<TYPE_X> xLocal = inQueueX.DeQue<TYPE_X>();
        Adds(yLocal, xLocal, (TYPE_X)0, length);
        outQueueY.EnQue<TYPE_Y>(yLocal);
        inQueueX.FreeTensor<TYPE_X>(xLocal);
    }

    __aicore__ inline void CopyOut(int i, int32_t progress, uint32_t length)
    {
        int32_t startIndex = i * totalLength + progress * tileLength;
        LocalTensor<TYPE_Y> yLocal = outQueueY.DeQue<TYPE_Y>();
        DataCopy(yGm[startIndex], yLocal, length);
        outQueueY.FreeTensor<TYPE_Y>(yLocal);
    }

    __aicore__ inline void ProcessV1()
    {
        int32_t loopCount = unitCount;
        for(int i=0; i<loopCount; i++)
        {
            for(int j=0; j<tileNum-1; j++)
            {
                CopyInV1(i, j, tileLength);
                Compute(i, j, tileLength);
                CopyOut(i, j, tileLength);
            }
                CopyInV1(i, tileNum - 1, lasttileLength);
                Compute(i, tileNum - 1, lasttileLength);
                CopyOut(i, tileNum - 1, lasttileLength);
        }
    }

    __aicore__ inline void CopyInV1(int i, int32_t progress, uint32_t length)
    {
        int32_t startIndex = GetOffset(i) + progress * tileLength + storageOffset;
        LocalTensor<TYPE_X> xLocal = inQueueX.AllocTensor<TYPE_X>();

        for(int j=0; j<length; j++)
        {
            xLocal.SetValue(j, ((__gm__ TYPE_X *)x_addr)[startIndex + j*stepSize]);   // cannot convert 'const AscendC::GlobalTensor<half>' to 'half' without a conversion operator
        }
        inQueueX.EnQue<TYPE_X>(xLocal);
    }


};


extern "C" __global__ __aicore__ void as_strided(GM_ADDR x, GM_ADDR size, GM_ADDR stride, GM_ADDR storage_offset, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    KernelAsStrided<DTYPE_X, DTYPE_X> op;
    op.Init(x, size, stride, storage_offset, y, tiling_data.blockLength, tiling_data.blockSize, tiling_data.ALIGN_NUM, tiling_data.Shape, tiling_data.dimNum, tiling_data.viewDim);
    op.Process();
}